import psycopg
from mylogger import logger
import os
from dotenv import load_dotenv
load_dotenv()

def init_database():
    table_name = 'text_v1'
    create_table_sql = f'''
    CREATE TABLE IF NOT EXISTS {table_name} (
        id SERIAL PRIMARY KEY,
        text TEXT NOT NULL,
        count INTEGER NOT NULL DEFAULT 0,
        llm_model VARCHAR(32) NOT NULL,
        version VARCHAR(32) NOT NULL,
        created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
    );
    CREATE INDEX IF NOT EXISTS index_{table_name}_1 ON {table_name}(version, llm_model, count);
    CREATE INDEX IF NOT EXISTS index_{table_name}_2 ON {table_name}(version, count);
    '''

    psql_url = os.environ.get('PSQL_URI')

    with psycopg.connect(psql_url) as conn:
        with conn.cursor() as cur:
            cur.execute(create_table_sql)
            # cur.execute("SELECT * FROM test")
            # cur.fetchone()
            # # will return (1, 100, "abc'def")

            # # You can use `cur.fetchmany()`, `cur.fetchall()` to return a list
            # # of several records, or even iterate on the cursor
            # for record in cur:
            #     print(record)
                
            conn.commit()
    logger.debug("init database successful")
            
            
def insert_data(text, count, llm_model, version):
    table_name = 'text_v1'
    insert_query = f"""
        INSERT INTO {table_name} (text, count, llm_model, version)
        VALUES (%s, %s, %s, %s);
    """

    psql_url = os.environ.get('PSQL_URI')

    with psycopg.connect(psql_url) as conn:
        with conn.cursor() as cur:
            cur.execute(insert_query, (text, count, llm_model, version))
            conn.commit()
    logger.debug("insert successful")
    
init_database()
